from os.path import join, abspath
from backend.experiment.framework.shots_prompt.task import \
    PerplexityPromptShotsClassificationTaskMaker
from backend.experiment.framework.task import TaskMaker
from backend.experiment.tokenizer import tokenizer
from backend.experiment.framework.shots_prompt.implement import \
    CommonShotsGenerator
from backend.experiment.model import context_limit
from backend.experiment.clue.afqmc.scheme0.implement \
    import AFQMCImplement as S0Impl


class AFQMCImplement(S0Impl):

    def get_task_maker(self) -> TaskMaker:
        return PerplexityPromptShotsClassificationTaskMaker(
            token_limit=context_limit,
            tokenizer=tokenizer,
            shots_generator=CommonShotsGenerator(
                length_sorted_example_tokens=self.get_sorted_example_tokens()
            ),
            label_map=self.label_map()
        )

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))
